module ModuleAGGREGATES_QTR
    
    use Globals
    use Toolbox
    
contains

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Measure Computation  w/out Qmu       !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    subroutine ComputeMEASURE_TR

    implicit none

    ! Local variable declaration
    real(8) :: muaux(NS)
    real(8) :: a_np, a_np_vec(1), da
    integer :: aind_np, aind_np_vec(1), is, is_np, ix_np, tt

    ! Transition first period measure: initial steady state
    mu_TR(:,1) = mu_init

    ! Iterate measure forward
    do tt = 1, T_TR-1

        muaux = 0d0     ! initialize updated measure

        do is = 1, NS

            ! linear interpolation in the asset dimension
            a_np        = min(a_pol_TR(is,tt),amax)
            a_np_vec    = a_np
            !aind_np_vec = vsearch_FET(a_np_vec,a_vec)
            aind_np_vec = vsearchprm_FET(a_np_vec,a_vec,aprm)
            aind_np     = aind_np_vec(1)
            da = (a_np - a_vec(aind_np))/(a_vec(aind_np+1) - a_vec(aind_np))

            ! account for all possible future productivity states
            do ix_np = 1, Nx
                is_np = aind_np + Na*(ix_np-1)
                muaux(is_np)   = muaux(is_np)   + (1.0D0-da)*Px(xind(is),ix_np)*mu_TR(is,tt)
                muaux(is_np+1) = muaux(is_np+1) +    da     *Px(xind(is),ix_np)*mu_TR(is,tt)
            end do

        enddo

        ! Store measure
        mu_TR(:,tt+1) = muaux

    enddo

    end subroutine ComputeMEASURE_TR

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!             Measure Aggregation         !!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  
    subroutine ComputeAGGREGATES_TR
    
    implicit none
    
    real(8), dimension(NS) :: yk, yl, ytot, transfers, et
    real(8) :: I0, I0_l, I0_k, I3, chit
    integer :: tt
    real(8), dimension(NS) :: muu, S1u, S2u 

    S1u = S(:,1)    ! assets
    S2u = S(:,2)    ! productivity

    !!! Initialize aggregates
    Kagg_TR     = 0.0D0
    Yagg_TR     = 0.0D0
    Lagg_TR     = 0.0D0
    Cagg_TR     = 0.0D0

    ! Loop over transition periods
    do tt = 1, T_TR
        
        muu = mu_TR(:,tt)       ! measure of that period
        
        Kagg_TR(tt)  = max(sum(S1u*muu) - Debt,1d-8)    ! agg. capital
        Cagg_TR(tt)  = sum(c_pol_TR(:,tt)*muu)
        Lagg_TR(tt)  = sum(S2u*h_pol_TR(:,tt)*muu)      ! agg. efficiency units of labor
        Yagg_TR(tt)  = (Kagg_TR(tt)**(1.0D0-alpha))*(Lagg_TR(tt)**alpha)                ! agg. output
        r_new_TR(tt) = - delta + (1.0D0-alpha)*((Lagg_TR(tt)/Kagg_TR(tt))**(alpha))     ! implied interest rate

        yk  = r_TR(tt)*S1u                  ! capital income
        yl  = wge_TR(tt)*S2u*h_pol_TR(:,tt) ! labor income

        ytot = yk+yl                        ! total income
        
        ! Tax revenue
        I0_l = sum(muu*exp(log(lambda_TR(tt))*((yl/ymean)**(-gamma)))*yl)   ! labor tax revenue
        I0_k = sum(muu*(tauk*yk))                                           ! capital tax revenue

        !  Transfer spending
        et   = exp(-rt*((yk+yl)/ymean-omegat))
        chit = exp(rt*omegat)/(1.0D0 + exp(rt*omegat))
        transfers = (mt_TR(tt)*ymean)*(et/(1.0D0+et))*(1D0/chit)
        I3 = sum(muu*(transfers))

        ! Government budget constraint
        ! spending + interest expenses + transfers - capital tax revenue - labor tax revenue - consumption tax revenue
        BC_TR(tt)  = (G + r_TR(tt)*Debt + I3 - I0_k - I0_l - tauc*Cagg_TR(tt))

        ! Asset market clearing
        ! private saving - govt debt - capital demand
        asset_TR(tt) = sum(S1u*muu) - Debt - ((r_TR(tt) + delta)/(1.0D0-alpha))**(-1.0D0/alpha) * Lagg_TR(tt)

    enddo

    end subroutine ComputeAGGREGATES_TR

       
end module ModuleAGGREGATES_QTR

